import numpy as np  # 导入NumPy库，用于数值计算
import faiss  # 导入FAISS库，用于高效的相似性搜索
from util import createXY  # 从util模块导入createXY函数，用于创建特征和标签
from sklearn.model_selection import train_test_split  # 导入train_test_split函数，用于划分数据集
from sklearn.neighbors import KNeighborsClassifier  # 导入KNeighborsClassifier类，用于K近邻分类
import argparse  # 导入argparse库，用于解析命令行参数
import logging  # 导入logging库，用于记录日志
from tqdm import tqdm  # 导入tqdm库，用于显示进度条
from FaissKNeighbors import FaissKNeighbors  # 从FaissKNeighbors模块导入FaissKNeighbors类
import os  # 导入os库，用于与操作系统交互
import sys  # 导入sys库，用于与Python解释器交互

# 配置日志记录的基本设置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 定义函数以解析命令行参数
def get_args():
    parser = argparse.ArgumentParser(description='使用CPU或GPU训练模型。')  # 创建ArgumentParser对象
    parser.add_argument('-m', '--mode', type=str, required=True, choices=['cpu', 'gpu'], help='选择训练模式：CPU或GPU。')  # 添加mode参数
    parser.add_argument('-f', '--feature', type=str, required=True, choices=['flat', 'vgg'], help='选择特征提取方法：flat或vgg。')  # 添加feature参数
    parser.add_argument('-l', '--library', type=str, required=True, choices=['sklearn', 'faiss'], help='选择使用的库：sklearn或faiss。')  # 添加library参数
    args = parser.parse_args()  # 解析命令行参数
    return args  # 返回解析后的参数

# 定义函数以检查CUDA环境
def check_cuda_environment():
    logging.info("检查 CUDA 环境...")  # 记录日志信息
    if 'CUDA_VISIBLE_DEVICES' in os.environ:  # 检查环境变量是否设置
        logging.info(f"CUDA_VISIBLE_DEVICES: {os.environ['CUDA_VISIBLE_DEVICES']}")  # 记录CUDA_VISIBLE_DEVICES的值
    else:
        logging.warning("未设置 CUDA_VISIBLE_DEVICES 环境变量")  # 记录警告信息

    try:
        import torch  # 尝试导入PyTorch库
        logging.info(f"PyTorch 版本: {torch.__version__}")  # 记录PyTorch版本
        logging.info(f"CUDA 是否可用: {torch.cuda.is_available()}")  # 记录CUDA是否可用
        if torch.cuda.is_available():  # 如果CUDA可用
            logging.info(f"CUDA 版本: {torch.version.cuda}")  # 记录CUDA版本
            logging.info(f"可用 GPU 数量: {torch.cuda.device_count()}")  # 记录可用GPU数量
    except ImportError:
        logging.warning("未安装 PyTorch，无法检查 CUDA 可用性")  # 记录警告信息

# 定义函数以检查GPU支持
def check_gpu_support():
    check_cuda_environment()  # 调用检查CUDA环境的函数
    try:
        logging.info(f"FAISS 库路径: {faiss.__file__}")  # 记录FAISS库路径
        logging.info(f"FAISS 编译信息: {faiss.get_compile_options()}")  # 记录FAISS编译信息
        gpu_count = faiss.get_num_gpus()  # 获取FAISS报告的GPU数量
        logging.info(f"FAISS 报告的 GPU 数量: {gpu_count}")  # 记录FAISS报告的GPU数量
        if gpu_count > 0:  # 如果有可用的GPU
            for i in range(gpu_count):  # 遍历每个GPU
                logging.info(f"测试 GPU {i}")  # 记录测试GPU的信息
                res = faiss.StandardGpuResources()  # 创建FAISS GPU资源
                logging.info(f"GPU {i} 测试成功")  # 记录GPU测试成功的信息
            logging.info(f"FAISS 检测到 {gpu_count} 个 GPU 设备")  # 记录检测到的GPU数量
            return True  # 返回True表示支持GPU
        else:
            logging.warning("FAISS 未检测到 GPU 设备，将使用 CPU 模式")  # 记录警告信息
            return False  # 返回False表示不支持GPU
    except AttributeError as e:
        logging.error(f"FAISS GPU 支持错误: {str(e)}")  # 记录错误信息
        logging.error("FAISS 不支持 GPU。请确保安装了支持 GPU 的 FAISS 版本")  # 记录错误信息
        return False  # 返回False表示不支持GPU
    except Exception as e:
        logging.error(f"检查 GPU 支持时发生未知错误: {str(e)}")  # 记录未知错误信息
        return False  # 返回False表示不支持GPU

# 定义主函数
def main():
    args = get_args()  # 获取命令行参数

    logging.info(f"Python 版本: {sys.version}")  # 记录Python版本
    logging.info(f"FAISS 版本: {faiss.__version__}")  # 记录FAISS版本

    # 检查GPU支持
    if args.mode == 'gpu':
        if not check_gpu_support():  # 如果不支持GPU
            logging.error("由于 GPU 不可用，程序将退出。请使用 -m cpu 选项重新运行程序")  # 记录错误信息
            return  # 退出程序

    try:
        res = faiss.StandardGpuResources() if args.mode == 'gpu' else None  # 创建FAISS GPU资源或设置为None
        logging.info("成功创建 GPU 资源" if args.mode == 'gpu' else "使用 CPU 模式")  # 记录创建GPU资源或使用CPU模式的信息
    except AttributeError:
        logging.error("创建 GPU 资源失败。可能的原因：")  # 记录错误信息
        logging.error("1. 安装的 FAISS 版本不支持 GPU")  # 记录可能的原因
        logging.error("2. CUDA 环境未正确配置")  # 记录可能的原因
        logging.error("请检查 FAISS 安装和 CUDA 配置，或使用 CPU 模式运行")  # 记录解决方案
        return  # 退出程序
    except Exception as e:
        logging.error(f"创建 GPU 资源时发生未知错误: {str(e)}")  # 记录未知错误信息
        return  # 退出程序

    logging.info(f"选择模式是 {args.mode.upper()}")  # 记录选择的模式
    logging.info(f"选择特征提取方法是 {args.feature.upper()}")  # 记录选择的特征提取方法
    logging.info(f"选择使用的库是 {args.library.upper()}")  # 记录选择的库

    try:
        X, y = createXY(train_folder="../data/train", dest_folder=".", method=args.feature)  # 创建特征和标签
        X = np.array(X).astype('float32')  # 将特征转换为NumPy数组并设置数据类型为float32
        faiss.normalize_L2(X)  # 对特征进行L2归一化
        y = np.array(y)  # 将标签转换为NumPy数组
        logging.info("数据加载和预处理完成。")  # 记录数据加载和预处理完成的信息
        logging.info(f"X 形状: {X.shape}, y 形状: {y.shape}")  # 记录特征和标签的形状
    except Exception as e:
        logging.error(f"数据加载或预处理失败：{str(e)}")  # 记录错误信息
        logging.error("请检查数据文件路径是否正确，以及数据格式是否符合要求")  # 记录解决方案
        return  # 退出程序

    try:
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)  # 划分数据集
        logging.info("数据集划分为训练集和测试集。")  # 记录数据集划分完成的信息
        logging.info(f"训练集形状: X_train {X_train.shape}, y_train {y_train.shape}")  # 记录训练集的形状
        logging.info(f"测试集形状: X_test {X_test.shape}, y_test {y_test.shape}")  # 记录测试集的形状
    except Exception as e:
        logging.error(f"数据集划分失败：{str(e)}")  # 记录错误信息
        return  # 退出程序

    best_k = -1  # 初始化最佳k值
    best_accuracy = 0.0  # 初始化最高准确率
    k_values = range(1, 6)  # 定义k值的范围

    KNNClass = FaissKNeighbors if args.library == 'faiss' else KNeighborsClassifier  # 根据选择的库设置KNN类
    logging.info(f"使用的库为: {args.library.upper()}")  # 记录使用的库

    try:
        for k in tqdm(k_values, desc='寻找最佳k值'):  # 遍历k值并显示进度条
            knn = KNNClass(k=k, res=res) if args.library == 'faiss' else KNNClass(n_neighbors=k)  # 创建KNN模型
            knn.fit(X_train, y_train)  # 训练KNN模型
            accuracy = knn.score(X_test, y_test)  # 评估模型准确率

            if accuracy > best_accuracy:  # 如果当前k值的准确率更高
                best_k = k  # 更新最佳k值
                best_accuracy = accuracy  # 更新最高准确率
    except Exception as e:
        logging.error(f"模型训练或评估过程中出错：{str(e)}")  # 记录错误信息
        if args.library == 'faiss' and args.mode == 'gpu':
            logging.error("如果使用 FAISS 和 GPU 模式，请确保 FAISS 正确安装并支持 GPU")  # 记录解决方案
        return  # 退出程序

    logging.info(f'最佳k值: {best_k}, 最高准确率: {best_accuracy}')  # 记录最佳k值和最高准确率

# 如果脚本作为主程序运行，则调用main函数
if __name__ == '__main__':
    main()